package org.example.sentinel;

import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;
import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class SentinelRuleManager {
    
    // 存储IP规则映射
    private final ConcurrentHashMap<String, FlowRule> ipRules = new ConcurrentHashMap<>();
    
    // 资源名称
    private static final String RESOURCE_NAME = "ipBasedControl";
    
    @PostConstruct
    public void init() {
        // 初始规则
        List<FlowRule> rules = new ArrayList<>();
        FlowRule defaultRule = new FlowRule();
        defaultRule.setResource(RESOURCE_NAME);
        defaultRule.setGrade(RuleConstant.FLOW_GRADE_QPS);
        defaultRule.setCount(100); // 默认QPS阈值
        rules.add(defaultRule);
        FlowRuleManager.loadRules(rules);
    }
    
    /**
     * 为异常IP添加流控规则
     */
    public Mono<Void> addRuleForIp(String ip) {
        return Mono.fromRunnable(() -> {
            List<FlowRule> rules = new ArrayList<>(FlowRuleManager.getRules());
            
            // 如果规则已存在，先移除
            if (ipRules.containsKey(ip)) {
                rules.remove(ipRules.get(ip));
            }
            
            // 创建针对特定IP的规则
            FlowRule ipRule = new FlowRule();
            ipRule.setResource(RESOURCE_NAME);
            ipRule.setGrade(RuleConstant.FLOW_GRADE_QPS);
            ipRule.setCount(0); // 完全限制访问
            ipRule.setLimitApp(ip); // 针对特定IP
            
            rules.add(ipRule);
            FlowRuleManager.loadRules(rules);
            ipRules.put(ip, ipRule);
        });
    }
    
    /**
     * 移除IP的流控规则
     */
    public Mono<Void> removeRuleForIp(String ip) {
        return Mono.fromRunnable(() -> {
            List<FlowRule> rules = new ArrayList<>(FlowRuleManager.getRules());
            if (ipRules.containsKey(ip)) {
                rules.remove(ipRules.get(ip));
                ipRules.remove(ip);
                FlowRuleManager.loadRules(rules);
            }
        });
    }
}